#include <fstream>
#include <dirent.h>
#include <vector>
#include <string.h>
#include <map>
#include <cmath>
#include <memory>
#include "acl/acl.h"
#include "acl/ops/acl_dvpp.h"
#define SHARED_PTR_DVPP_BUF(buf) (shared_ptr<uint8_t>((uint8_t *)(buf), [](uint8_t* p) { acldvppFree(p); }))
#define INFO_LOG(fmt, ...) fprintf(stdout, "[INFO]  " fmt "\n", ##__VA_ARGS__);fflush(stdout)
#define ERROR_LOG(fmt, ...)fprintf(stderr, "[ERROR] " fmt "\n", ##__VA_ARGS__)
#define YUV420SP_SIZE(width, height) ((width) * (height) * 3 / 2)
#define ALIGN_UP(num, align) (((num) + (align) - 1) & ~((align) - 1))
#define ALIGN_UP2(num) ALIGN_UP(num, 2)
#define ALIGN_UP16(num) ALIGN_UP(num, 16)
#define ALIGN_UP64(num) ALIGN_UP(num, 64)

using namespace std;
typedef enum Result {
    SUCCESS = 0,
    FAILED = 1
} Result;

struct ImageData {
    uint32_t width = 0;
    uint32_t height = 0;
    uint32_t alignWidth = 0;
    uint32_t alignHeight = 0;
    uint32_t size = 0;
    std::shared_ptr<uint8_t> data;
};

class SampleDVPP {
    public:
    SampleDVPP(int32_t device);
    ~SampleDVPP();
    Result InitResource();
    Result ReadBinFile(const string imageFile,ImageData &image);
    Result JpegDecode(ImageData imageJpg,ImageData &imageYuv);
    Result Resize(ImageData imageYuv,ImageData &imageVpc,int width, int height);
    Result SaveBinFile(ImageData image,string name);
    private:
    void ReleaseResource();
    int32_t deviceId_;
    aclrtStream stream_;
    acldvppChannelDesc *dvppChannelDesc_;
};

SampleDVPP::SampleDVPP(int32_t device) :deviceId_(device)
{
}

// 析构函数里释放资源
SampleDVPP::~SampleDVPP()
{
    ReleaseResource();
}

// AscendCL资源初始化
Result SampleDVPP::InitResource() {
    
    // init acl resource
    const char *aclConfigPath = "";
    aclError ret = aclInit(aclConfigPath);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("aclInit failed, errorCode is %d", ret);
        return FAILED;
    }
    // 设置device
    ret = aclrtSetDevice(deviceId_);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("aclrtSetDevice failed, errorCode is %d", ret);
        return FAILED;
    }
    // 创建stream
    ret = aclrtCreateStream(&stream_);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("aclrtCreateStream failed, errorCode is %d", ret);
        return FAILED;
    }
    // 创建dvpp通道
    dvppChannelDesc_ = acldvppCreateChannelDesc();
    ret = acldvppCreateChannel(dvppChannelDesc_);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("acldvppCreateChannel failed, errorCode is %d", ret);
        return FAILED;
    }

}

// 将图片读取到HOST内存中
Result SampleDVPP::ReadBinFile(const string imageFile,ImageData &image)
{
    std::ifstream binFile(imageFile, std::ifstream::binary);
    if(!binFile.good()) {
        ERROR_LOG("dog.jpg does not exit!");
        return FAILED;
    }
    uint32_t offset = 0;
    binFile.seekg(offset, binFile.end);
    uint32_t binFileBufferLen = binFile.tellg();
    binFile.seekg(offset, binFile.beg);
    uint8_t* binFileBufferData = new(std::nothrow) uint8_t[binFileBufferLen];
    binFile.read((char *)binFileBufferData, binFileBufferLen);
    binFile.close();
    image.data.reset(binFileBufferData, [](uint8_t* p) { delete[](p); });
    image.size = binFileBufferLen;
    return SUCCESS;
}

// 进行JpegD解码,将jpg转换为YUV420格式图片
Result SampleDVPP::JpegDecode(ImageData imageJpg,ImageData &imageYuv)
{
    // 将图片HOST内存拷贝到DEVICE
    void *deviceJpgBuffer;
    acldvppMalloc(&deviceJpgBuffer, imageJpg.size);
    
    aclError ret = aclrtMemcpy(deviceJpgBuffer, imageJpg.size, imageJpg.data.get(),
                               imageJpg.size, ACL_MEMCPY_HOST_TO_DEVICE);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("copy image dvpp failed, errorCode is %d", ret);
        return FAILED;
    }
    // 获取原始图片宽高,内存大小信息
    int32_t components = 0;
    acldvppJpegGetImageInfo(imageJpg.data.get(), imageJpg.size,
                            &(imageJpg.width), &(imageJpg.height),
                            &components);
    acldvppPicDesc *decodeOutputDesc = acldvppCreatePicDesc(); 
    // alignment with different chip version
    //auto socVersion = aclrtGetSocName();
    // 计算输出图片的宽高对齐  
    imageYuv.width = ALIGN_UP2(imageJpg.width);
    imageYuv.height = ALIGN_UP2(imageJpg.height);
    imageYuv.alignWidth = ALIGN_UP64(imageJpg.width); // 64-byte alignment
    imageYuv.alignHeight = ALIGN_UP16(imageJpg.height); // 16-byte alignment
    // 输出图片YUV420格式内存大小计算为  (widthStride * heightStride) * 3 / 2
    imageYuv.size = YUV420SP_SIZE(imageYuv.alignWidth, imageYuv.alignHeight);

    void *deviceYuvBuffer;
    acldvppMalloc(&deviceYuvBuffer, imageYuv.size);
    imageYuv.data = SHARED_PTR_DVPP_BUF(deviceYuvBuffer);
    // 创建输出图片描述信息并设置属性
    acldvppSetPicDescData(decodeOutputDesc, imageYuv.data.get());
    acldvppSetPicDescFormat(decodeOutputDesc, PIXEL_FORMAT_YUV_SEMIPLANAR_420);
    acldvppSetPicDescWidth(decodeOutputDesc, imageYuv.width);
    acldvppSetPicDescHeight(decodeOutputDesc, imageYuv.height);
    acldvppSetPicDescWidthStride(decodeOutputDesc, imageYuv.alignWidth);
    acldvppSetPicDescHeightStride(decodeOutputDesc, imageYuv.alignHeight);
    acldvppSetPicDescSize(decodeOutputDesc, imageYuv.size);

    // JpegD解码异步接口执行
    acldvppJpegDecodeAsync(dvppChannelDesc_, deviceJpgBuffer,
                           imageYuv.size, decodeOutputDesc, stream_);
    // 阻塞等待流执行完成
    ret = aclrtSynchronizeStream(stream_);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("acldvppJpegDecodeAsync failed, errorCode is %d", ret);
        return FAILED;
    }

    //资源释放
    acldvppDestroyPicDesc(decodeOutputDesc);
    acldvppFree(deviceJpgBuffer);
    
}

Result SampleDVPP::Resize(ImageData imageYuv,ImageData &imageVpc,int width, int height)
{
    aclError ret;
    acldvppPicDesc *vpcInputDesc = acldvppCreatePicDesc();
    acldvppPicDesc *vpcOutputDesc = acldvppCreatePicDesc();
    acldvppResizeConfig *resizeConfig = acldvppCreateResizeConfig();

    // 设置VPC输入图片描述信息
    acldvppSetPicDescData(vpcInputDesc, imageYuv.data.get());  // set input desc
    acldvppSetPicDescFormat(vpcInputDesc, PIXEL_FORMAT_YUV_SEMIPLANAR_420);
    acldvppSetPicDescWidth(vpcInputDesc, imageYuv.width);
    acldvppSetPicDescHeight(vpcInputDesc, imageYuv.height);
    acldvppSetPicDescWidthStride(vpcInputDesc, imageYuv.alignWidth);
    acldvppSetPicDescHeightStride(vpcInputDesc, imageYuv.alignHeight);
    acldvppSetPicDescSize(vpcInputDesc, imageYuv.size);

    // 计算输出图片宽高对齐
    imageVpc.width = ALIGN_UP2(width);
    imageVpc.height = ALIGN_UP2(height);
    imageVpc.alignWidth = ALIGN_UP16(imageVpc.width);
    imageVpc.alignHeight = ALIGN_UP2(imageVpc.height);
    imageVpc.size = YUV420SP_SIZE(imageVpc.alignWidth, imageVpc.alignHeight);
    void *vpcOutBufferDev; // vpc output buffer
    acldvppMalloc(&vpcOutBufferDev, imageVpc.size);
    imageVpc.data = SHARED_PTR_DVPP_BUF(vpcOutBufferDev);
    // 设置输出图片描述信息
    acldvppSetPicDescData(vpcOutputDesc, imageVpc.data.get());   // set output desc
    acldvppSetPicDescFormat(vpcOutputDesc, PIXEL_FORMAT_YUV_SEMIPLANAR_420);
    acldvppSetPicDescWidth(vpcOutputDesc, imageVpc.width);
    acldvppSetPicDescHeight(vpcOutputDesc, imageVpc.height);
    acldvppSetPicDescWidthStride(vpcOutputDesc, imageVpc.alignWidth);
    acldvppSetPicDescHeightStride(vpcOutputDesc, imageVpc.alignHeight);
    acldvppSetPicDescSize(vpcOutputDesc, imageVpc.size);

    // VPC Resize异步接口执行
    acldvppVpcResizeAsync(dvppChannelDesc_, vpcInputDesc,
                          vpcOutputDesc, resizeConfig, stream_);
    // 阻塞等待流执行结束
    ret = aclrtSynchronizeStream(stream_);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("acldvppVpcResizeAsync failed, errorCode is %d", ret);
        return FAILED;
    }
    // 资源释放
    acldvppDestroyPicDesc(vpcInputDesc);
    acldvppDestroyPicDesc(vpcOutputDesc);
    acldvppDestroyResizeConfig(resizeConfig);

}

// 资源释放
void SampleDVPP::ReleaseResource()
{
    aclError ret;
    ret = acldvppDestroyChannel(dvppChannelDesc_);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("destroy channel failed, errorCode is %d", ret);
    }
    ret = acldvppDestroyChannelDesc(dvppChannelDesc_);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("destroy channel description failed, errorCode is %d", ret);
    }
    if (stream_ != nullptr) {
        ret = aclrtDestroyStream(stream_);
        if (ret != ACL_SUCCESS) {
            ERROR_LOG("aclrtDestroyStream failed, errorCode is %d", ret);
        }
        stream_ = nullptr;
    }
    ret = aclrtResetDevice(deviceId_);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("aclrtResetDevice failed, errorCode is %d", ret);
    }

    ret = aclFinalize();
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("aclFinalize failed, errorCode is %d", ret);
    }
}


Result SampleDVPP::SaveBinFile(ImageData image,string name) {
    void *binFile;
    aclrtMallocHost(&binFile,image.size);
    aclrtMemcpy(binFile, image.size, image.data.get(),
                               image.size, ACL_MEMCPY_DEVICE_TO_HOST);
    std::ofstream file(name, std::ios::out | std::ios::binary);
    if (!file) {
        ERROR_LOG("Failed to open file for writing");
        return FAILED;
    }
    
    file.write(static_cast<char*>(binFile), image.size);
    
    if (!file) {
        ERROR_LOG("Failed to write data to file");
        return FAILED;
    }
    
    file.close();
    aclrtFreeHost(binFile);
    return SUCCESS;
}

int main()
{
    const string imagePath = "../data/dog.jpg";
    int32_t device = 0;
    SampleDVPP sampleDVPP(device);

    Result ret = sampleDVPP.InitResource();
    if (ret != SUCCESS) {
        ERROR_LOG("InitResource failed");
        return FAILED;
    }

    ImageData imageJpg;
    ret = sampleDVPP.ReadBinFile(imagePath,imageJpg);
    if(ret != SUCCESS)  {
        ERROR_LOG("load input picture failed!");
        return FAILED;
    }
    
    ImageData imageYuv;
    ret = sampleDVPP.JpegDecode(imageJpg,imageYuv);
    if(ret != SUCCESS)  {
        ERROR_LOG("jpeg decode failed!");
        return FAILED;
    }
    ImageData imageVpc;
    ret = sampleDVPP.Resize(imageYuv,imageVpc,224,224);
    if(ret != SUCCESS)  {
        ERROR_LOG("resize failed!");
        return FAILED;
    }
    // 保存输出文件到本地
    sampleDVPP.SaveBinFile(imageJpg,"jpg.yuv");
    sampleDVPP.SaveBinFile(imageVpc,"resize.yuv");
}